//! 对称加密
use aes::Aes256;
use block_modes::block_padding::Pkcs7;
use block_modes::{BlockMode, Cbc};
use sha2::{Digest, Sha256};

type Aes256CbcType = Cbc<Aes256, Pkcs7>;
/// Aes256-Cbc对称加密
///
/// key 32个字节
///
/// iv  16个字节
///
/// # Examples
///
/// ```ignore
/// fn main(){
///     let source = b"hellow word";
///     let key = "123457890abcdef";
///     let iv = "56789abcdef";
///     let aes = Aes256Cbc::new(key, iv).unwrap();
///     let enc = aes.encrypt(source);
///     let dec = aes.decrypt(&enc).unwrap();
///     assert_eq!(source.as_ref(), dec.as_slice());
/// }
/// ```
pub struct Aes256Cbc {
    key: [u8; 32],
    iv: [u8; 16],
}
impl Aes256Cbc {
    pub fn new(key: &str, iv: &str) -> Result<Self, String> {
        let (k, v) = Self::sha2(key, iv);
        Ok(Aes256Cbc { key: k, iv: v })
    }
    /// 更新key和iv
    pub fn update_key_iv(&mut self, key: &str, iv: &str) {
        let (k, v) = Self::sha2(key, iv);
        self.key = k;
        self.iv = v;
    }
    pub fn encrypt(&self, d: &[u8]) -> Vec<u8> {
        Aes256CbcType::new_var(self.key.as_ref(), self.iv.as_ref())
            .unwrap()
            .encrypt_vec(d)
    }

    pub fn decrypt(&self, d: &[u8]) -> Result<Vec<u8>, String> {
        Aes256CbcType::new_var(&self.key[..], &self.iv[..])
            .unwrap()
            .decrypt_vec(d)
            .map_err(|e| e.to_string())
    }
    fn sha2(key: &str, iv: &str) -> ([u8; 32], [u8; 16]) {
        let mut hasher = Sha256::new();
        let mut hasher2 = hasher.clone();
        hasher.input(key);
        let k = hasher.result();
        let mut _key = [0; 32];
        let mut _iv = [0; 16];
        _key.copy_from_slice(k.as_ref());
        hasher2.input(iv);
        let i = hasher2.result();
        _iv.copy_from_slice(&i[10..26]);
        (_key, _iv)
    }
}

#[test]
fn aes256_cbc_encrypt_decrypt() {
    let source = b"hellow word";
    let key = "123457890abcdef";
    let iv = "56789abcdef";
    let aes = Aes256Cbc::new(key, iv).unwrap();
    let enc = aes.encrypt(source);
    let dec = aes.decrypt(&enc).unwrap();
    assert_eq!(source.as_ref(), dec.as_slice());
}
